877F - Ann and Books - CodeForces Solution


data structures flows hashing *2300

Please click on ads to support us..

C++ Code:

#include "bits/stdc++.h"
#include "ext/pb_ds/assoc_container.hpp"
//#include "ext/pb_ds/tree_policy.hpp"
using namespace __gnu_pbds;
using namespace std;
inline void read() {}
template<typename T, typename ...Ts> inline void read(T& x, Ts&... vals) {
    x = 0;
    int c = getchar();
    while (isspace(c)) c = getchar();
    bool neg = (c == '-');
    for (neg ? c = getchar() : c; '0' <= c && c <= '9'; c = getchar()) x = (x * 10) + (c - '0');
    if (neg) x = -x;
    read(vals...);
}
template<class T1, class T2> ostream& operator << (ostream& ostr, pair<T1, T2>& p) {
    ostr << '(' << p.first << ' ' << p.second << ')';
    return ostr;
}
template<class T> ostream& operator << (ostream& ostr, vector<T>& vc) {
    ostr << '(';
    for (int i = 0; i < vc.size(); i++) ostr << vc[i] << (i != vc.size()-1 ? " " : "");
    ostr << ')';
    return ostr;
}
template<class T> ostream& operator << (ostream& ostr, set<T>& st) {
    for (auto &e: st) ostr << e << ' ';
    return ostr;
}
template<class T1, class T2> ostream& operator << (ostream& ostr, map<T1, T2>& mp) {
    for (auto& [x, y] : mp) ostr << '[' << x << " -> " << y << "] ";
    return ostr;
}
#define dbg(...) cerr << "[" << #__VA_ARGS__ << "] : " << __VA_ARGS__ << endl

typedef long long ll;
//typedef tree<ll, null_type, less<>, rb_tree_tag, tree_order_statistics_node_update> ordered_set;
//#define read(k) cin >> k

struct query {
    int idx, l, r;
    ll ans;
};

void run_test() {
    int n, k, q;
    read(n, k);
    vector<int> v(n), a(n);
    vector<ll> pre(n+1), comp = {-k, 0, k};
    for (int& e: v) read(e);
    for (int i = 1; i <= n; i++) {
        read(a[i-1]);
        pre[i] = pre[i-1] + (v[i-1] == 1 ? a[i-1] : -a[i-1]);
        comp.push_back(pre[i]);
        comp.push_back(pre[i]-k);
        comp.push_back(pre[i]+k);
    }
    sort(comp.begin(), comp.end());
    gp_hash_table<ll, int, hash<ll>> com_val;
    int ptr = 0;
    com_val[comp[0]] = ptr++;
    for (int i = 1; i < comp.size(); i++) {
        if (comp[i] != comp[i-1]) com_val[comp[i]] = ptr++;
    }
    vector<array<int, 3>> cor(n+1);
    for (int i = 0; i <= n; i++) {
        cor[i][0] = com_val[pre[i]-k];
        cor[i][1] = com_val[pre[i]];
        cor[i][2] = com_val[pre[i]+k];
    }
    read(q);
    vector<query> qq(q);
    for (int i = 0; i < q; i++) {
        qq[i].idx = i;
        read(qq[i].l, qq[i].r);
        --qq[i].l, --qq[i].r;
    }
    int BLOCK_SIZE = (int)sqrt(n);
    sort(qq.begin(), qq.end(), [&BLOCK_SIZE](query& q1, query& q2) {
        int x = q1.l / BLOCK_SIZE, y = q2.l / BLOCK_SIZE;
        if (x != y) return q1.l < q2.l;
        else if (x&1) return q1.r > q2.r;
        else return q1.r < q2.r;
    });
    ll cur_ans = 0;
    vector<int> mp1(ptr), mp2(ptr);
    auto add = [&] (int idx, int left) {
        mp1[cor[idx][1]]++, mp2[cor[idx+1][1]]++;
        if (left) cur_ans += mp2[cor[idx][2]];
        else cur_ans += mp1[cor[idx+1][0]];
    };
    auto remove = [&] (int idx, bool left) {
        if (left) cur_ans -= mp2[cor[idx][2]];
        else cur_ans -= mp1[cor[idx+1][0]];
        mp1[cor[idx][1]]--, mp2[cor[idx+1][1]]--;
    };
    int L = qq[0].l, R = qq[0].r;
    for (int p = L; p <= R; p++) add(p, false);
    qq[0].ans = cur_ans;
    for (int i = 1; i < q; i++) {
        for (int p = L-1; p >= qq[i].l; p--) add(p, true);
        for (int p = R+1; p <= qq[i].r; p++) add(p, false);
        for (int p = L; p < qq[i].l; p++) remove(p, true);
        for (int p = R; p > qq[i].r; p--) remove(p, false);
        qq[i].ans = cur_ans;
        L = qq[i].l, R = qq[i].r;
    }
    sort(qq.begin(), qq.end(), [](query& q1, query& q2) {
        return q1.idx < q2.idx;
    });
    for (auto& e: qq) cout << e.ans << '\n';
}

signed main() {
//    ::freopen("/media/sda4/CODES/CXX/AOPS/input.txt", "r", stdin);
//    ::freopen("/media/sda4/CODES/CXX/AOPS/output.txt", "w", stdout);
    cin.tie(nullptr)->sync_with_stdio(false);
    int t = 1;
//    read(t);
    while (t--) run_test();
    return 0;
}


Comments

Submit
0 Comments
More Questions

866. Prime Palindrome
1516A - Tit for Tat
622. Design Circular Queue
814. Binary Tree Pruning
791. Custom Sort String
787. Cheapest Flights Within K Stops
779. K-th Symbol in Grammar
701. Insert into a Binary Search Tree
429. N-ary Tree Level Order Traversal
739. Daily Temperatures
647. Palindromic Substrings
583. Delete Operation for Two Strings
518. Coin Change 2
516. Longest Palindromic Subsequence
468. Validate IP Address
450. Delete Node in a BST
445. Add Two Numbers II
442. Find All Duplicates in an Array
437. Path Sum III
436. Find Right Interval
435. Non-overlapping Intervals
406. Queue Reconstruction by Height
380. Insert Delete GetRandom O(1)
332. Reconstruct Itinerary
368. Largest Divisible Subset
377. Combination Sum IV
322. Coin Change
307. Range Sum Query - Mutable
287. Find the Duplicate Number
279. Perfect Squares